"""
经验回放缓冲区

用于存储和管理强化学习轨迹数据。
"""

from __future__ import annotations

import numpy as np
import torch


class RolloutBuffer:
    """
    经验回放缓冲区
    
    存储轨迹数据，包括观测、动作、奖励、价值估计、对数概率和终止标志。
    支持高效的数据存储和批量获取。
    """
    
    def __init__(self):
        """初始化缓冲区"""
        self.observations = []
        self.actions = []
        self.rewards = []
        self.values = []
        self.log_probs = []
        self.dones = []
    
    def add(
        self,
        obs: np.ndarray,
        action: np.ndarray,
        reward: float,
        value: float,
        log_prob: float,
        done: bool,
    ):
        """
        添加一步经验
        
        Args:
            obs: 观测 [obs_dim]
            action: 动作 [action_dim]
            reward: 奖励（标量）
            value: 价值估计（标量）
            log_prob: 对数概率（标量）
            done: 是否终止（布尔值）
        """
        self.observations.append(obs)
        self.actions.append(action)
        self.rewards.append(reward)
        self.values.append(value)
        self.log_probs.append(log_prob)
        self.dones.append(done)
    
    def clear(self):
        """清空缓冲区"""
        self.observations.clear()
        self.actions.clear()
        self.rewards.clear()
        self.values.clear()
        self.log_probs.clear()
        self.dones.clear()
    
    def get(self) -> tuple[torch.Tensor, ...]:
        """
        获取所有经验并转换为张量
        
        Returns:
            observations: [T, obs_dim]
            actions: [T, action_dim]
            rewards: [T]
            values: [T]
            log_probs: [T]
            dones: [T]
        """
        observations = torch.FloatTensor(np.array(self.observations))
        actions = torch.FloatTensor(np.array(self.actions))
        rewards = torch.FloatTensor(self.rewards)
        values = torch.FloatTensor(self.values)
        log_probs = torch.FloatTensor(self.log_probs)
        dones = torch.FloatTensor(self.dones)
        
        return observations, actions, rewards, values, log_probs, dones
    
    def __len__(self) -> int:
        """返回缓冲区中的步数"""
        return len(self.rewards)
    
    def is_empty(self) -> bool:
        """检查缓冲区是否为空"""
        return len(self) == 0
